#!/usr/bin/env python
"""
简单的启动脚本，用于导入三元组数据到Neo4j数据库
"""

import os
import sys
import logging
from dotenv import load_dotenv
from import_triples_to_neo4j import TripleImporter

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler("import.log"),
        logging.StreamHandler()
    ]
)

def main():
    """主函数"""
    # 加载环境变量
    load_dotenv()
    
    # 获取Neo4j连接信息
    neo4j_uri = os.getenv("NEO4J_URI", "bolt://localhost:7687")
    neo4j_user = os.getenv("NEO4J_USER", "neo4j")
    neo4j_password = os.getenv("NEO4J_PASSWORD", "password")
    
    # 获取文件路径
    if len(sys.argv) > 1:
        file_path = sys.argv[1]
    else:
        file_path = "12k_all.txt"  # 默认使用12k_all.txt文件
    
    if not os.path.exists(file_path):
        print(f"错误: 文件 {file_path} 不存在")
        return
    
    # 创建导入器
    importer = TripleImporter(neo4j_uri, neo4j_user, neo4j_password, batch_size=1000)
    
    try:
        # 创建索引
        print("创建Neo4j索引...")
        importer.create_indexes()
        
        # 导入数据
        print(f"开始导入文件: {file_path}")
        stats = importer.import_triples(file_path)
        
        # 输出结果
        if stats.get("success", False):
            print("\n======== 导入结果 ========")
            print(f"成功导入 {stats['total_triples']} 个三元组")
            print(f"包含 {stats['unique_entities']} 个唯一实体")
            print(f"包含 {stats['unique_relationships']} 种关系类型")
            
            print("\n实体类型分布:")
            for entity_type, examples in stats['entity_types'].items():
                print(f"  {entity_type}: {len(examples)} 个样例: {', '.join(examples)}")
                
            print("\n导入已完成! 数据已成功导入到Neo4j数据库。")
            print("现在您可以通过Neo4j Browser访问数据，例如尝试以下查询:")
            print("  MATCH (n) RETURN n LIMIT 100;  // 显示100个节点")
            print("  MATCH (n)-[r]->(m) RETURN n, r, m LIMIT 50;  // 显示50个关系")
            print("  MATCH (n:Malware) RETURN n;  // 显示所有恶意软件")
        else:
            print(f"导入失败: {stats.get('error', '未知错误')}")
    finally:
        importer.close()

if __name__ == "__main__":
    main() 